from ultralytics import YOLO
import os
from datetime import datetime

# 加载模型
model = YOLO('models/yolov8n-seg.pt')

# 输入图片路径
image_path = r'H:\xiaomi\test4\original_rtsp_frame_486_vehicles1_20251119_133922.jpg'

# 创建带时间戳的输出目录
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
output_dir = f'results/{timestamp}'
os.makedirs(output_dir, exist_ok=True)

print("开始预测...")
print(f"输入图片: {image_path}")
print(f"输出目录: {output_dir}")

# 进行预测并保存结果
results = model.predict(
    source=image_path,
    save=True,           # 保存带预测框的图片
    save_txt=True,       # 保存检测框数据
    save_conf=True,      # 保存置信度
    project=output_dir,  # 输出目录
    name='predict',      # 子目录名称
    exist_ok=True        # 允许覆盖已存在目录
)

print("预测完成！")
print(f"结果保存在: {output_dir}")

# 打印检测信息
for i, result in enumerate(results):
    print(f"\n图片 {i+1}:")
    if result.boxes is not None:
        print(f"检测到 {len(result.boxes)} 个目标")
        for box in result.boxes:
            cls_name = result.names[int(box.cls[0])]
            conf = box.conf[0].item()
            print(f"  - {cls_name}: {conf:.3f}")